import pandas as pd
import matplotlib.pyplot as plt
from tqdm import tqdm

from utils import search_data


gongye_data = pd.read_excel(f"cache/{search_data('生产测_工业')}", 
    index_col=0)['工业增加值当月同比（合并1-2月）']

fuwuye_data = pd.read_excel(f"cache/{search_data('生产测_服务业')}", 
    index_col=0)['服务业生产指数：当月同比']

df_supply = pd.concat([gongye_data, fuwuye_data], axis=1)
df_supply.columns = ['工业增加值当月同比', '服务业生产指数']
df_supply.index = pd.to_datetime(df_supply.index)

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

import os
from datetime import datetime

# 绘制多指标柱状图
def plot_multi_bar(df, title='多指标柱状图', figsize=(14, 8), save_path=None, 
                  colors=None, show_values=False, add_grid=True):
    """
    绘制美化的多指标柱状图
    
    参数:
    df: DataFrame, 包含多个指标的数据
    title: str, 图表标题
    figsize: tuple, 图表大小
    save_path: str, 保存路径，None表示自动生成
    colors: list, 柱状图颜色列表，None表示使用默认颜色
    show_values: bool, 是否显示数值标签
    add_grid: bool, 是否添加网格线
    """
    # 创建图形
    fig, ax = plt.subplots(figsize=figsize)
    
    # 设置柱状图宽度
    width = 0.8 / len(df.columns)
    
    # 获取数据
    index = df.index
    columns = df.columns
    
    # 默认颜色列表
    default_colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b']
    if colors is None:
        colors = default_colors[:len(columns)]
    
    # 绘制柱状图
    bars = []
    for i, (column, color) in enumerate(zip(columns, colors)):
        x = [j + i * width for j in range(len(index))]
        bar = ax.bar(x, df[column], width, label=column, color=color, alpha=0.85)
        bars.append(bar)
        
        # 显示数值标签
        if show_values:
            for rect in bar:
                height = rect.get_height()
                ax.annotate(f'{height:.1f}',
                            xy=(rect.get_x() + rect.get_width() / 2, height),
                            xytext=(0, 3),  # 3点垂直偏移
                            textcoords="offset points",
                            ha='center', va='bottom', fontsize=9)
    
    # 设置x轴标签
    ax.set_xticks([j + width * (len(columns) - 1) / 2 for j in range(len(index))])
    ax.set_xticklabels([date.strftime('%Y-%m') for date in index], rotation=45, ha='right')
    
    # 添加标题和标签（美化样式）
    ax.set_title(title, fontsize=16, fontweight='bold', pad=20)
    ax.set_xlabel('日期', fontsize=12, labelpad=10)
    ax.set_ylabel('同比增长率(%)', fontsize=12, labelpad=10)
    
    # 添加图例（美化位置）
    ax.legend(loc='best', fontsize=11, frameon=True, framealpha=0.9)
    
    # 添加网格线
    if add_grid:
        ax.grid(axis='y', linestyle='--', alpha=0.7)
    
    # 美化图表边框
    for spine in ['top', 'right']:
        ax.spines[spine].set_visible(False)
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图表
    if save_path:
        # 确保保存目录存在
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"图表已保存至: {save_path}")
    else:
        # 自动生成保存路径
        today = datetime.now().strftime('%Y-%m-%d')
        auto_dir = f"plot/plot_{today}"
        os.makedirs(auto_dir, exist_ok=True)
        auto_save_path = f"{auto_dir}/工业与服务业对比图.png"
        plt.savefig(auto_save_path, dpi=300, bbox_inches='tight')
        print(f"图表已自动保存至: {auto_save_path}")
    
    # 显示图表
    plt.show()
    plt.close()

# 使用函数绘制df_supply的多指标柱状图（带美化效果）

plot_multi_bar(
    df_supply.iloc[-20:,:], 
    title='工业增加值与服务业生产指数同比增长率对比分析',
    colors=['#4E79A7', '#F28E2B'],  # 蓝色和橙色
    show_values=True,  # 显示数值标签
    add_grid=True  # 添加网格线
)